package org.dslforge.aurum;

import static org.dslforge.aurum.Symbol.*;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.Stack;

class AugmentedGrammar {
	final List<Symbol> nonterminals = new ArrayList<Symbol>();

	final List<Symbol> terminals = new ArrayList<Symbol>();

	final List<Production> productions = new ArrayList<Production>();

	final Map<Symbol, Collection<Symbol>> firstSets = new HashMap<Symbol, Collection<Symbol>>();

	final Map<Symbol, Collection<Symbol>> followSets = new HashMap<Symbol, Collection<Symbol>>();

	final GrammarDefinition definition;

	final Production startProduction;

	private final Symbol start = nonterminal("$start");

	AugmentedGrammar(GrammarDefinition definition, Symbol start) {
		this.definition = definition;
		startProduction = new Production(this.start, start);
		productions.add(startProduction);
		discoverTerminalsAndNonterminals(start);
		computeFistSets();
		computeFollowSets();
	}

	private void discoverTerminalsAndNonterminals(Symbol start) {
		Stack<Symbol> unvisited = new Stack<Symbol>();
		Set<Symbol> visited = new HashSet<Symbol>();
		unvisited.push(start);
		nonterminals.add(start);
		visited.add(start);
		while (!unvisited.isEmpty()) {
			Symbol visiting = unvisited.pop();
			for (Production production : definition.productions.get(visiting)) {
				productions.add(production);
				for (Symbol symbol : production.symbols)
					if (!visited.contains(symbol)) {
						if (symbol.isNonterminal) {
							unvisited.push(symbol);
							nonterminals.add(symbol);
						} else
							terminals.add(symbol);
						visited.add(symbol);
					}
			}
		}
	}

	private void computeFistSets() {
		for (Symbol terminal : terminals)
			getFirstSet(terminal).add(terminal);
		List<Production> productions = new ArrayList<Production>();
		for (Production production : this.productions)
			if (production.leftEnd().isNonterminal)
				productions.add(production);
			else
				getFirstSet(production.nonterminal).add(production.leftEnd());
		boolean changed = true;
		while (changed) {
			changed = false;
			for (Production production : productions)
				for (Symbol symbol : production.symbols) {
					changed = addAllExceptEpsilon(getFirstSet(symbol),
							getFirstSet(production.nonterminal));
					if (!firstSets.get(symbol).contains(EPSILON))
						break;
				}
		}
	}

	private void computeFollowSets() {
		getFollowSet(start).add(EOF);
		List<Production> productions = new ArrayList<Production>();
		for (Production production : this.productions) {
			if (production.rightEnd().isNonterminal)
				productions.add(production);
			for (int i = 0; i < production.symbols.length - 1; i++)
				if (production.get(i).isNonterminal)
					addAllExceptEpsilon(getFirstSet(production.get(i + 1)), getFollowSet(production
							.get(i)));
		}
		boolean changed = true;
		while (changed) {
			changed = false;
			for (Production production : productions) {
				Collection<Symbol> follow = getFollowSet(production.nonterminal);
				changed |= getFollowSet(production.rightEnd()).addAll(follow);
				for (int i = production.symbols.length - 1; i > 0; i--)
					if (getFirstSet(production.get(i)).contains(EPSILON)
							&& production.get(i - 1).isNonterminal)
						changed |= getFollowSet(production.get(i - 1)).addAll(follow);
					else
						break;
			}
		}
	}

	private boolean addAllExceptEpsilon(Collection<Symbol> from, Collection<Symbol> to) {
		boolean changed = false;
		for (Symbol terminal : from)
			if (!terminal.equals(EPSILON))
				changed |= to.add(terminal);
		return changed;
	}

	Collection<Symbol> getFirstSet(Symbol symbol) {
		return getSet(firstSets, symbol);
	}

	Collection<Symbol> getFollowSet(Symbol nonterminal) {
		return getSet(followSets, nonterminal);
	}

	private <T> Collection<Symbol> getSet(Map<T, Collection<Symbol>> sets, T key) {
		Collection<Symbol> set = sets.get(key);
		if (set == null) {
			set = new HashSet<Symbol>();
			sets.put(key, set);
		}
		return set;
	}

	Collection<LR0Item> closureForLR0Items(Collection<LR0Item> items) {
		Set<LR0Item> closure = new HashSet<LR0Item>();
		Stack<LR0Item> unmarked = new Stack<LR0Item>();
		for (LR0Item item : items) {
			closure.add(item);
			unmarked.push(item);
		}
		while (!unmarked.isEmpty()) {
			LR0Item visiting = unmarked.pop();
			if (!visiting.isHandle() && visiting.expected.isNonterminal)
				for (Production production : definition.productions.get(visiting.expected)) {
					LR0Item item = new LR0Item(production, 0);
					if (!closure.contains(item)) {
						closure.add(item);
						unmarked.push(item);
					}
				}
		}
		return closure;
	}

	Collection<LR0Item> goToForLR0Items(Collection<LR0Item> items, Symbol symbol) {
		Set<LR0Item> goTo = new HashSet<LR0Item>();
		for (LR0Item item : items)
			if (!item.isHandle() && item.expected.equals(symbol))
				goTo.add(new LR0Item(item.production, item.position + 1));
		return closureForLR0Items(goTo);
	}

	Collection<LR1Item> closureForLR1Items(Collection<LR1Item> items) {
		Set<LR1Item> closure = new HashSet<LR1Item>();
		Stack<LR1Item> unmarked = new Stack<LR1Item>();
		for (LR1Item item : items) {
			closure.add(item);
			unmarked.push(item);
		}
		while (!unmarked.isEmpty()) {
			LR1Item visiting = unmarked.pop();
			LR0Item item = visiting.item;
			if (!item.isHandle() && item.expected.isNonterminal) {
				Set<Symbol> lookaheads = new HashSet<Symbol>();
				Symbol beta = item.production.get(item.position + 1);
				if (beta != null)
					lookaheads.addAll(firstSets.get(beta));
				if (lookaheads.size() == 0 || lookaheads.contains(EPSILON))
					lookaheads.add(visiting.lookahead);
				for (Production production : definition.productions.get(item.expected))
					for (Symbol lookahead : lookaheads) {
						LR1Item newItem = new LR1Item(production, 0, lookahead);
						if (!closure.contains(newItem)) {
							closure.add(newItem);
							unmarked.push(newItem);
						}
					}
			}
		}
		return closure;
	}

	Collection<LR1Item> goToForLR1Items(Collection<LR1Item> items, Symbol symbol) {
		Collection<LR1Item> goTo = new HashSet<LR1Item>();
		for (LR1Item LR1 : items)
			if (!LR1.item.isHandle() && LR1.item.expected.equals(symbol))
				goTo.add(new LR1Item(LR1.item.production, LR1.item.position + 1, LR1.lookahead));
		return closureForLR1Items(goTo);
	}

	List<Collection<LR0Item>> itemsCollection() {
		List<Collection<LR0Item>> collection = new ArrayList<Collection<LR0Item>>();
		Stack<Collection<LR0Item>> workingList = new Stack<Collection<LR0Item>>();
		Collection<LR0Item> I0 = closureForLR0Items(Collections.singleton(new LR0Item(
				startProduction, 0)));
		collection.add(I0);
		workingList.push(I0);
		while (!workingList.isEmpty()) {
			Collection<LR0Item> visiting = workingList.pop();
			for (Symbol symbol : expectedSymbols(visiting)) {
				Collection<LR0Item> newItems = goToForLR0Items(visiting, symbol);
				if (!collection.contains(newItems)) {
					collection.add(newItems);
					workingList.push(newItems);
				}
			}
		}
		return collection;
	}

	private Set<Symbol> expectedSymbols(Collection<LR0Item> items) {
		Set<Symbol> symbols = new HashSet<Symbol>();
		for (LR0Item item : items)
			if (!item.isHandle())
				symbols.add(item.expected);
		return symbols;
	}

	boolean isKernel(LR0Item item) {
		return item.position != 0 || item.production.equals(startProduction);
	}
}